#include "libavutil/opt.h"
#include "libavutil/pixdesc.h"
#include "libavutil/avassert.h"
#include "libavutil/intreadwrite.h"
#include "avcodec.h"
#include "internal.h"
#include "parser.h"
#include <float.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include "rmH265.h"

#if 1
/*************************************************************************************/

static int calculate_nalsize(unsigned char* src,int nal_length_size)
{
	int nalsize,i;
	nalsize = 0;
	for(i = 0; i < nal_length_size; i++)
	{
		nalsize = (nalsize << 8) |src[i];
	}
	return nalsize;
}
static int avc_stream_change(unsigned char* src,int size,int nal_length_size,unsigned char* dst)
{
	int offset = 0;
	while( size > 0 )
	{
		int nalsize,i;
		nalsize = calculate_nalsize(src,nal_length_size);
		if( nalsize+nal_length_size > size || nalsize <= 0 )
		{
			//LOGI("err,nalsize=%d %d\r\n",nalsize,size);
			//return -1;
			break;
		}
		//WriteLog2("nalsize=%d %d",nalsize,size);
		dst[0]=0;dst[1]=0;dst[2]=0;dst[3]=1;
		dst += 4;
		offset += (nalsize+4);
		memcpy(dst,src+nal_length_size ,nalsize);
		dst += nalsize;
		
		size -= (nalsize+nal_length_size);		
		src += (nalsize+nal_length_size);
	}
	return offset;
}
#define MAX_FRAME_BUFFER 8
#define MAX_PTS 32
typedef struct X265DecContext {
    AVClass        *class;  // need by opt.c
    void           *dec;
    ParseContext    parse_context;
    AVFrame         frame;

    int flags;        ///< AVCodecContext.flags (HQ, MV4, ...)
    int flags2;       ///< AVCodecContext.flags2

	RM_VIDEO_CODECAPI          rmhevcdecApis;
	RM_HANDLE                   h265dec;
	char* extradata;
	char* configdata;
	int configdata_size;
	int change_flag;
	int injectconfig;
	/////
	int in_frames;
	char* frame_buffer[MAX_FRAME_BUFFER];
	int64_t frame_pts[MAX_FRAME_BUFFER];
	int wpos;
	int rpos;
	int fullflag;
	//pts buffer 
//   	int64_t frame_pts[MAX_PTS];
//	int frame_pts_writepos;
//	int frame_pts_readpos;		
} X265DecContext;

static int ff_h265_decode_extradata(X265DecContext *s, const uint8_t *buf, int size)
{
    if (!buf || size <= 0)
        return -1;
	uint8_t * dst = (uint8_t*)s->configdata;
	uint8_t * p = buf;
	uint8_t * pend = buf+size-1;
	while( p < pend )
	{
		if( (p[0] == 0) && (p[1] == 1) )
		{
			p += 2;
			int n = AV_RB16(p);			
			p+=2;
			
			dst[0] = 0;dst[1] = 0;dst[2] = 0;dst[3] = 1;
			dst += 4;
			memcpy(dst,p,n);
			p += n;
			dst += n;
		} else		
			++p;
	}
	s->configdata_size = (int)(dst-(uint8_t*)s->configdata);
	return 0;
}
static av_cold int X265Dec_init(AVCodecContext *avctx)
{
	RM_S32 nRet;
    X265DecContext *s = avctx->priv_data;

    avctx->pix_fmt = PIX_FMT_YUV420P;
    avcodec_get_frame_defaults(&s->frame);
    s->frame.width  = -1;
    s->frame.height = -1;

	nRet = rmGetDecHandle(&s->rmhevcdecApis);
	s->h265dec = NULL;
	
	if(nRet != RM_RETURN_OK){
		av_log(s, AV_LOG_DEBUG,"err,rmGetDecHandle %d\r\n",nRet);	
		return AVERROR(ENOMEM);
	}
	
	nRet = s->rmhevcdecApis.Open(&s->h265dec);
	if(nRet != RM_RETURN_OK){
		av_log(s, AV_LOG_DEBUG,"err,rmhevcdecApis.Open %d\r\n",nRet);	
		return AVERROR(ENOMEM);
	}
	s->extradata = av_malloc(1024*512);
	s->configdata = av_malloc(128);
	int cpu_num = 4;
 	nRet = s->rmhevcdecApis.Set(s->h265dec, RM_PID_CPUNUM, (RM_PTR)&cpu_num); 
	av_log(s, AV_LOG_DEBUG,"set thread count=%d,ret=%d\n",cpu_num,nRet);
    if (avctx->extradata_size > 0 && avctx->extradata )
    {
    	ff_h265_decode_extradata(s,avctx->extradata,avctx->extradata_size);
		av_log(s, AV_LOG_DEBUG,"configdata_size=%d\n",s->configdata_size);
		s->change_flag = 1;
	#if 0
		{
			FILE*fp=fopen("/sdcard/extra","wb");
			if(fp)
			{
				fwrite(s->configdata,1, s->configdata_size,fp);
				fclose(fp);
			}
		}
	#endif
    }else
	    s->change_flag = -1;
	av_log(s, AV_LOG_DEBUG,"s->h265dec=%p",s->h265dec);
    return 0;
}
static av_cold int X265Dec_close(AVCodecContext *avctx)
{
	int i;
    X265DecContext *s = avctx->priv_data;	
	av_log(s, AV_LOG_DEBUG,"265Decoder,close");
	if( s->h265dec )
		s->rmhevcdecApis.Close(s->h265dec);
	for(i=0;i<MAX_FRAME_BUFFER;++i)
	{
		if( s->frame_buffer[i] )
			av_freep(&s->frame_buffer[i]);
	}
	if( s->extradata )
		av_freep(&s->extradata);
	if( s->configdata )
		av_freep(&s->configdata);
//    if (s->frame.data[0])
//        av_free(s->frame.data[0]);
    return 0;
}
static void Init_Frame_Buffer(AVCodecContext *avctx,int w,int h)
{
	int i;
    X265DecContext *s  = avctx->priv_data;
	//av_log(s, AV_LOG_DEBUG,"w=%d,h=%d",w,h);
	if( w != avctx->width )
		avctx->width = w;
	if( h != avctx->height )
		avctx->height = h;		
    if (w != s->frame.width || h != s->frame.height) 
	{
		av_log(s, AV_LOG_DEBUG,"265Decoder,trace(%d %d) frame %p",w,h,s->frame.data[0]);
        if (s->frame.data[0])
            av_free(s->frame.data[0]);
        s->frame.data[0] = av_malloc(w * h * 3 / 2);
//		av_log(s, AV_LOG_DEBUG,"trace frame %p",s->frame.data[0]);
        s->frame.data[1] = s->frame.data[0] + w * h;
        s->frame.data[2] = s->frame.data[1] + w * h / 4;

        s->frame.linesize[0] = w;
        s->frame.linesize[1] = w/2;
        s->frame.linesize[2] = w/2;
        s->frame.width  = w;
        s->frame.height = h;
		s->frame.format = 0;
		#if 0
		for(i=0;i<MAX_FRAME_BUFFER;++i)
		{
			if( s->frame_buffer[i] )
				av_free(s->frame_buffer[i]);
			s->frame_buffer[i] = av_malloc(w * h * 3 / 2);
			//av_log(s, AV_LOG_DEBUG,"265Decoder,frame_buffer[%d]=%p\n",i,s->frame_buffer[i]);
		}
		#endif
    }
}
static int GetBufferSize(X265DecContext *s)
{
	if( s->fullflag )
		return MAX_FRAME_BUFFER;
	
	if( s->wpos >= s->rpos )
		return s->wpos-s->rpos;
	else
		return MAX_FRAME_BUFFER-(s->rpos-s->wpos);
}

//pull all decoded yuv frames from decoder
static void Pull_AllDecoderFrames(AVCodecContext *avctx)
{
    X265DecContext *s  = avctx->priv_data;
	RM_U32                  nRet;
	RM_VIDEO_BUFFER         out_data;
	RM_VIDEO_BASICINFO    pStreamInfo;
	int timeout = 0;	
retry:	
	nRet = s->rmhevcdecApis.Process(s->h265dec, (RM_PTR)&out_data, &pStreamInfo);
	//av_log(s, AV_LOG_DEBUG,"s->rmhevcdecApis.Process=%x %d",nRet,buf_size);
	if(nRet == RM_RETURN_OK)
	{
		//get a yuv frame
		if(pStreamInfo.Format.Type != RM_VIDEO_FRAME_NULL)
		{
			int w = pStreamInfo.Format.Width;
			int h = pStreamInfo.Format.Height;
			Init_Frame_Buffer(avctx,w,h);

			if( GetBufferSize(s) >= MAX_FRAME_BUFFER )
			{
				av_log(s, AV_LOG_DEBUG,"265Decoder,buffer full lose frame");
			} else
			{				
				int i;
				char* dst = s->frame_buffer[s->wpos];
				if( NULL == dst )
				{
					s->frame_buffer[s->wpos] = av_malloc(w * h * 3 / 2);
					dst = s->frame_buffer[s->wpos];
				}
				s->frame_pts[s->wpos] = out_data.Time;
				//y
				uint8_t* src = out_data.Buffer[0];
				for(i=0;i<h;++i)
				{
					memcpy(dst,src,w);
					dst += w;
					src += out_data.Stride[0];
				}
				//u
				src = out_data.Buffer[1];
				for(i=0;i<h/2;++i)
				{
					memcpy(dst,src,w/2);
					dst += w/2;
					src += out_data.Stride[1];
				}
				//v
				src = out_data.Buffer[2];
				for(i=0;i<h/2;++i)
				{
					memcpy(dst,src,w/2);
					dst += w/2;
					src += out_data.Stride[2];
				}
				++s->wpos;
				if( s->wpos >= MAX_FRAME_BUFFER )
					s->wpos = 0;
				if( s->wpos == s->rpos )
					s->fullflag = 1;
//				av_log(s, AV_LOG_DEBUG,"265Decoder,get frame wpos=%d,rpos=%d",s->wpos,s->rpos);
			}
			if( pStreamInfo.Flag )
				goto retry;
		}
		else
		{			
			if( GetBufferSize(s) > 0 )
				return;
			//wait decoder some time
			if( s->in_frames > 4 )
			{
				timeout++;
				usleep(1000*2);
				if( timeout < 20 )
					goto retry;
			}
			av_log(s, AV_LOG_DEBUG,"265Decoder[%d]timeout=%d,Format.Type=%x(used=%d)",s->in_frames,timeout,pStreamInfo.Format.Type,pStreamInfo.UsedLens);
		}
	}
	else //if (nRet != RM_RETURN_INPUT_NO_ENOUGH)
	{
		if( GetBufferSize(s) <= 0 )
			av_log(s, AV_LOG_DEBUG,"265Decoder[%d],err=%x(used=%d)",s->in_frames,nRet,pStreamInfo.UsedLens);
	}	
}
static int X265Dec_frame(AVCodecContext *avctx, void *data, int *data_size, AVPacket *avpkt)
{
	RM_CODECBUFFER          in_data;
    X265DecContext *s  = avctx->priv_data;
    const uint8_t *buf = avpkt->data;
    int buf_size       = avpkt->size;
    AVFrame *pict      = data;

	if( data_size )
	    *data_size = 0;
    s->flags  = avctx->flags;
    s->flags2 = avctx->flags2;    
    /* no supplementary picture */
    if (buf_size == 0) {
        /* special case for last picture */
		//av_log(s, AV_LOG_DEBUG,"265Decoder,hint,len=%d",avpkt->size);
        //return avpkt->size;
        goto checkbuffer;
    }
#if 0	
	s->frame_pts[s->frame_pts_writepos] = avpkt->pts;
	s->frame_pts_writepos++;
	if( s->frame_pts_writepos >= MAX_PTS )
		s->frame_pts_writepos = 0;
#endif	
//av_log(s, AV_LOG_DEBUG,"265Decoder trace %d(%x %x %x %x)(%x %x %x %x)\r\n",avpkt->size,avpkt->data[0],avpkt->data[1],avpkt->data[2],avpkt->data[3]
//		,avpkt->data[4],avpkt->data[5],avpkt->data[6],avpkt->data[7]);
#if 1
	if( -1 == s->change_flag )
	{		
		if( ((buf[0]==0)&&(buf[1]==0)&&(buf[2]==0)&&(buf[3]==1))||((buf[0]==0)&&(buf[1]==0)&&(buf[2]==1)) )
			s->change_flag = 0;
		else
		{
			int n = calculate_nalsize(buf,4);
			if( n+4 == buf_size )
			{
				s->change_flag = 1;
			}
			av_log(s, AV_LOG_DEBUG,"reset change_flag %d %d",n,buf_size);
		}
	}
#endif

	if( 0 == s->injectconfig )
	{
		s->injectconfig = 1;
		if( (s->configdata_size > 0) && (s->configdata!=NULL) )
		{
			in_data.Buffer = s->configdata;
			in_data.Length = s->configdata_size;
			s->rmhevcdecApis.Set(s->h265dec, RM_PID_INPUTDATA, (RM_PTR)&in_data);	
		}
	}
	if( s->change_flag )
	{
		in_data.Length = avc_stream_change(buf,buf_size,4,(unsigned char*)s->extradata);		
		in_data.Buffer = s->extradata;
		//av_log(s, AV_LOG_DEBUG,"trace change,%d,%d",in_data.Length,buf_size);
	}else
	{
		in_data.Buffer = buf;
		in_data.Length = buf_size;
	}
	#if 0
		{
			static FILE*fp=NULL;
			if( NULL == fp ) 
				{
				fp = fopen("/sdcard/huawei.265","wb");
				av_log(s, AV_LOG_DEBUG,"fp---=%p",fp);
				}
			if(fp)
			{
				fwrite(in_data.Buffer,1,in_data.Length,fp);
				fflush(fp);
				av_log(s, AV_LOG_DEBUG,"write frame %d",in_data.Length);
			}
			return avpkt->size;
		}
	#endif
	s->in_frames++;	
	in_data.Time = avpkt->pts;
//	av_log(s, AV_LOG_DEBUG,"trace settime,pts=%d dts=%d size=%d",(int)in_data.Time,(int)avpkt->dts,in_data.Length);
	s->rmhevcdecApis.Set(s->h265dec, RM_PID_INPUTDATA, (RM_PTR)&in_data);	
	Pull_AllDecoderFrames(avctx);
checkbuffer:	
	if( GetBufferSize(s) > 0 )
	{
#if 0	
		s->frame.pts = s->frame_pts[s->frame_pts_readpos];
		s->frame_pts_readpos++;
		if( s->frame_pts_readpos >= MAX_PTS )
			s->frame_pts_readpos = 0;
#endif		
		s->frame.pts = s->frame_pts[s->rpos];
		memcpy(s->frame.data[0],s->frame_buffer[s->rpos],s->frame.width*s->frame.height*3/2);
		if( data_size )
	        *data_size = sizeof(AVFrame);
		if( pict )
	        *pict      = s->frame;		
		s->rpos++;
		if( s->rpos >= MAX_FRAME_BUFFER )
			s->rpos = 0;
		if( s->fullflag )
			s->fullflag = 0;
//		av_log(s, AV_LOG_DEBUG,"265Decoder,out frame wpos=%d,rpos=%d pts=%lld,in=%lld",s->wpos,s->rpos,s->frame.pts,avpkt->pts);
	}
	return avpkt->size;
}
static void X265Dec_flush(AVCodecContext *avctx)
{
    X265DecContext *s = avctx->priv_data;
	av_log(s, AV_LOG_DEBUG,"265Decoder,flush");
	s->rmhevcdecApis.Set(s->h265dec, RM_PID_FLUSH, NULL);
	s->in_frames = 0;
	s->rpos = s->wpos = 0;
	s->fullflag = 0;
//	s->frame_pts_readpos = s->frame_pts_writepos = 0;
}

AVCodec ff_h265_decoder = {
    .name           = "hevc",
    .type           = AVMEDIA_TYPE_VIDEO,
    .id             = AV_CODEC_ID_H265,
    .priv_data_size = sizeof(X265DecContext),
    .init           = X265Dec_init,
    .close          = X265Dec_close,
    .decode         = X265Dec_frame,
    .capabilities   = CODEC_CAP_DR1,
    .flush          = X265Dec_flush,
    .long_name        = NULL_IF_CONFIG_SMALL("h265 Decoder for HEVC / H.265"),
    .profiles       = NULL,
};
#endif
